from ultralytics import YOLO
import cv2
import os

# 类别索引到名称的映射（用于验证，实际写入的是索引）
CLASS_MAPPING = {
    0: 'hel',
    1: 'nohel',
    2: 'wheel'
}

def process_folder(folder_path, output_dir=None):
    """处理整个文件夹的图片，生成YOLO格式标签文件"""
    # 支持的图片格式
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']

    # 加载YOLOv8模型
    model = YOLO(r'D:\CodeCNN\yolov8-study\runs\detect\train30\weights\best30.pt')

    # 遍历文件夹
    for file in os.listdir(folder_path):
        file_path = os.path.join(folder_path, file)

        # 检查是否为支持的图片文件
        if os.path.isfile(file_path) and any(file.lower().endswith(ext) for ext in image_extensions):
            print(f"Processing: {file_path}")

            # 读取图片获取尺寸
            img = cv2.imread(file_path)
            if img is None:
                print(f"  Failed to read image: {file_path}")
                continue

            height, width, _ = img.shape

            # 进行预测
            results = model.predict(source=file_path, conf=0.25)

            # 存储检测结果（用于写入文件）
            detections = []
            count_stats = {name: 0 for name in CLASS_MAPPING.values()}  # 统计用

            for result in results:
                for box in result.boxes:
                    cls_idx = int(box.cls.item())
                    if cls_idx not in CLASS_MAPPING:
                        continue  # 跳过非目标类别

                    # 获取边界框 (xmin, ymin, xmax, ymax)
                    xyxy = box.xyxy[0].tolist()

                    xmin, ymin, xmax, ymax = xyxy
                    # 转换为YOLO格式：归一化中心坐标 + 宽高
                    x_center = ((xmin + xmax) / 2) / width
                    y_center = ((ymin + ymax) / 2) / height
                    bbox_width = (xmax - xmin) / width
                    bbox_height = (ymax - ymin) / height

                    # 写入类别索引（不是类别名）
                    class_name = CLASS_MAPPING[cls_idx]
                    detections.append(f"{cls_idx} {x_center:.6f} {y_center:.6f} {bbox_width:.6f} {bbox_height:.6f}")

                    # 统计数量
                    count_stats[class_name] += 1

            # 设置输出路径
            if output_dir is None:
                output_dir = folder_path
            os.makedirs(output_dir, exist_ok=True)

            # 保存YOLO格式的.txt标签文件
            txt_filename = os.path.splitext(file)[0] + ".txt"
            txt_path = os.path.join(output_dir, txt_filename)

            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write("\n".join(detections))

            print(f"  Saved YOLO label to: {txt_path}")
            print(f"  Detections: {count_stats}")

    print("\nBatch processing completed!")


if __name__ == "__main__":
    # 配置路径
    input_folder = r"D:\Datasets\pic0910-公司未戴头盔软件\images"   # 图像输入路径
    output_folder = r"D:\Datasets\pic0910-公司未戴头盔软件\labels"  # .txt标签输出路径

    # 处理整个文件夹
    process_folder(input_folder, output_folder)